#!/usr/bin/perl

use strict;
use Getopt::Long;

# read all biallelic genotypes from per chromosome files within this directory
my $MDAdir = "/Net/spud/raid/AFFY/MouseDiversityArray/";
my $CCsamplefile = "/Net/spud/raid/CC-Cousin-Lines-SNP6/CClineIDs.csv";

my $inputdir = "/Net/spud/raid/CC-Cousin-Lines-SNP6/ANALYSIS-070809/genotype.19122009";
my $error = 0.001;
my $prefix = "preCC.IL";
my $regex="Mott|REDO";

my @founders = ("A/J", "C57BL/6J", "129S1SvImJ", "NOD/LtJ", "NZO/HILtJ", "CAST/EiJ", "PWK/PhJ", "WSB/EiJ");
my @founders = ( "129S1SvImJ", "A/J", "AKR/J",  "BALB/cJ", "C3H/HeJ", "C57BL/6J", "LP/J", "NOD/LtJ", "NZO/HILtJ");
my $snpmap = "/Net/spud/raid/CC-Cousin-Lines-SNP6/snpmap.txt";
my $monomorphic = 1;

GetOptions(  "error=s"=>\$error, "dir=s"=>\$inputdir, "regex=s"=>\$regex, "prefix=s"=>\$prefix, "mono=s"=>\$monomorphic );
warn "inputdir $inputdir\n";


#_________________________________________________________________________________________
#
#  Read sample file data
#
#_________________________________________________________________________________________
my $sampleinfo;
open(CC,$CCsamplefile)|| die "Could not open $CCsamplefile\n";
$_ = <CC>;
while (<CC>) {
  chomp;
  my ( $Position,$run,$CCline,$Generation,$Sex,$subline,$coatcolor,$run,$ID) = split(/,/);
  $sampleinfo->{$ID} = [$Position,$run,$CCline,$Generation,$Sex,$subline,$coatcolor,$run,$ID];
}







my ($inbredgeno, $ac) = ReadInbredGenotypes( $MDAdir, \@founders);
my ( $snpinfo, $genotype, $igenotype, $map, $jaxmap ) = ReadSNPMap($snpmap);
my $perlegen;

#my $perlegen = ReadPerlegen($perlegendir, $map);

foreach my $chr ( sort keys %$map ) {
  my @sorted = sort { $a->[1] <=> $b->[1] } @{$map->{$chr}};
  $map->{$chr} = \@sorted;
}



#_________________________________________________________________________________________
#
#  Read Genotype data
#
#           probeset_id   CC04REDO.CEL CC10REDO.CEL  CC43REDO.CEL  CC65REDO.CEL CC72REDO.CEL ...
#           JAX00000001   2            0             -1            1            0            ...
#
#_________________________________________________________________________________________

warn "inputdir $inputdir\n";
my $genofile = "$inputdir/quant-norm.pm-only.brlmm-p.calls.txt";
open(GENO,$genofile) || die "Could not open $genofile\n";
print "reading $genofile\n";
my $n=0;
my @samples;
#  my $inbred;
my $use;
my $geno;
my $freq;

while (<GENO>) 
{
	chomp;
	if ( ! /^\#/ ) 
	{   

		# 
		# 	Read lines from first (header) line, and discard those with names not matching
		# 		regular expression
		# 		
		if ( $n++ == 0 ) 
		{      
			my $dummy;
			($dummy,@samples) = split;
			for (my $k=0;$k<=$#samples;$k++) 
			{	#	  $inbred->[$k] = 1 if ( $samples[$k] =~ /Inbred/ );
				if ( $samples[$k] =~ /$regex/i ) 
				{	$use->[$k] = 1;
					print "using $samples[$k]\n";
				}
				else 
				{	  
					print "skipping $samples[$k]\n";
				}	
			}
		} else 
		# 
		# 	Read data: probset_id (JAX00000XXX) and genotypes as 0 / -1 /2
		# 	
		#	only use matching probes ids and animals with matching ids (from header)
		# 		
		{      
			my ( $jax, @g ) = split;
			for (my $k=0;$k<=$#g;$k++) 
			{	
				if ( $use->[$k] && $inbredgeno->{$jax}) 
				{	  
					$geno->{$samples[$k]}{$jax} = $genotype->{$jax}[$g[$k]+1];
					$freq->{$jax}{$geno->{$samples[$k]}{$jax}}++;
				}
			}	
		}
	}
}

my $maf;
my $m=0;
my $r=0;
my $mono;
foreach my $jax ( keys %$freq ) {
  my @k = keys %{$freq->{$jax}};
  my $w=$k[0];
  my $t=0;
  foreach my  $k (@k ){
    $t += $freq->{$jax}{$k};
    $w = $k if ( $freq->{$jax}{$k} <= $freq->{$jax}{$w} )
  }
  $maf->{$jax} = $freq->{$jax}{$w}; 
  $m += $mono->{$jax} = $maf->{$jax}==$t || $maf->{$jax}==1;
  print "$jax $maf->{$jax} $freq->{$jax}{$w} $t\n" if ( $r++ <1000);
}

print "$m monomorphic snps\n";

if ( 0 ) {
  my $match;
  print "matching inbred lines\n";
  foreach my $jax ( sort keys %$inbredgeno ) {
    my $ref = $jaxmap->{$jax};
    my $igeno = $inbredgeno->{$jax};
    if ( exists( $perlegen->{$ref->[0]}{$ref->[1]} )) {
      my $pg = $perlegen->{$ref->[0]}{$ref->[1]};
      foreach my $inbred ( keys %$pg ) {
	if ( defined($pg->{$inbred})) {
	  foreach my $sample ( keys %$igeno ) {
	    $match->{$sample}{$inbred} ++ if ( defined($igeno->{$sample}) && $pg->{$inbred} eq $igeno->{$sample});
	  }
	}
      }
    }
  }
  my $max;
  my $best;
  my $match2;
  foreach my $sample (sort keys %$match ) {
    $max->{$sample} =0;
    foreach my $inbred ( sort keys %{$match->{$sample}} ) {
      $match2->{$inbred}{$sample} = $match->{$sample}{$inbred};
      if ( $max->{$sample} < $match->{$sample}{$inbred} ) {
	$max->{$sample} = $match->{$sample}{$inbred};
	$best->{$sample} = $inbred;
      }
    }
    print "$sample -> $best->{$sample} $max->{$sample}\n";
  }
  
  my $max2;
  my $best2;
  
  foreach my $inbred ( sort keys %$match2) {
    $max2->{$inbred} =0;
    foreach my $sample (sort keys %{$match2->{$inbred}} ) {
      if ( $max2->{$inbred} < $match2->{$inbred}{$sample} ) {
	$max2->{$inbred} = $match2->{$inbred}{$sample};
	$best2->{$inbred} = $sample;
      }
    }
    print "$inbred -> $best2->{$inbred} $max2->{$inbred}\n";
  }
}

	
foreach my $chr ( sort keys %$map ) {
  my $pedfile = "$chr.$prefix.ped";
  open(PED,">$pedfile") || die "Could not open $pedfile\n";
  print "writing $pedfile\n";

  foreach my $sample ( sort keys %$geno ) {
    my $id = $sample;
    my $sex = "M";
    my $fam = "IL";
    if ( $sampleinfo->{$id} ) {
      $id = $sampleinfo->{$id}[2] . "-" . $sampleinfo->{$id}[5];
      $sex = $sampleinfo->{$id}[4];
      $fam = $sampleinfo->{$id}[2];
    }	
    print PED "CC $id 0 0 M NA ";
    my $ns=0;
    my $na=0;
    foreach my $ref ( @{$map->{$chr}} ) {
      my $jax = $ref->[0];
      if ( $inbredgeno->{$jax} && !$mono->{$jax}) {
	print PED $geno->{$sample}{$jax}? $geno->{$sample}{$jax} : " NA NA";
	if ( $geno->{$sample}{$jax} ) {
	  $ns++;
	} else {
	  $na++;
	}
      }	
    }	
    print PED "\n";
    my $sum = $ns+$na;
    print "chr $sample $ns $na $sum\n";
  }
  close(PED);

  my $mapfile = "$chr.$prefix.map";
  open(MAP,">$mapfile") || die "Could not open $mapfile\n";
  print MAP "marker\tchromosome\tbp\n";
  my $allelesfile = "$chr.$prefix.alleles";
  open(ALLELES,">$allelesfile") || die "Could not open $allelesfile\n";
  print "writing $allelesfile\n";
  my $nm = 0;
  foreach my $ref ( @{$map->{$chr}} ) {
    $nm ++ if ( $inbredgeno->{$ref->[0]} && !$mono->{$ref->[0]} );
  }	
  print ALLELES "markers $nm strains 8\nstrain_names " . join( " ", @founders) . "\n";
  foreach my $ref ( @{$map->{$chr}} ) {
    my $jax = $ref->[0];
    if ( $inbredgeno->{$jax} && !$mono->{$jax} ) {
      print MAP join( "\t", $jax, $chr, $ref->[1]) . "\n";
      my $cm = $ref->[1]/2.0e6;
      print ALLELES "marker $jax 3 $chr $cm\n";
      foreach my $a ( sort keys %{$ac->{$jax}} ) {
	print ALLELES "allele $a ";
	my $p = $ac->{$jax}{$a}/8;
	foreach my $s ( @founders ) {
	  printf ALLELES " %.4f", $error+ $p*($inbredgeno->{$jax}{$s} eq $a);
	}	
	print ALLELES "\n";
      }	
      print ALLELES "allele NA ";
      foreach my $s ( @founders ) {
	printf ALLELES " %.4f", 1/8;
      }	
      print ALLELES "\n";
    }	
  }
  close(MAP);
  close(ALLELES);
}

sub ReadPerlegen{

  my ( $perlegendir, $map ) = @_;
  my $trans = { "AA" => "A", "CC" => "C", "GG" => "G", "TT" => "T","A" => "A", "C" => "C", "G" => "G", "T" => "T", "NN" => "NA", "0" => "NA", "NA"=>"NA"  };

  my $geno;
  my @strains;
  my $want;
  my $mat=0;
  foreach my $chr ( keys %$map ) {
    my $m=0;
    foreach my $ref ( @{$map->{$chr}} ) {
      $want->{$chr}{$ref->[1]} = $ref->[0];
    }
  }
  
  foreach my $chr ( sort {$a <=> $b} keys %$want ) {
    my $c = $chr;
    my ($c) = $chr =~ /chr(.+)/;
    if ( $c > 0 ) {
      $c = sprintf( "Chr%2.2d", $c);
    } else {
      $c = "Chr$c";
    }
    my $perlfile = "$perlegendir/b04_$c\_genotype.dat";
    if ( open(PERLEGEN, $perlfile) ) {
      $_ = <PERLEGEN>;
      chomp;
      s/ /./g;
      my ($local, $id, $chrom, $acc, $bp, $s, $al);
      ($local, $id, $chrom, $acc, $bp, $s, $al, @strains ) = split(/\t/);
      while (<PERLEGEN>) {
	my ($local, $id, $chrom, $acc, $bp, $s, $al, @g ) = split(/\t/);
	if ( $want->{$chr}{$bp} ) {
	  $mat++;
	  for (my $s=0;$s<=$#strains;$s++ ) {
	    $geno->{$chr}{$bp}{$strains[$s]} = $g[$s] ne "N" ? $g[$s] : undef;
	  }
	}
      }
      close(PERLEGEN);
    }	
  }
  print "mat = $mat\n";
  return($geno);
}

sub ReadSNPMap {
  my ($snpmap) = @_;
  my $snpinfo;
  my $genotype;
  my $igenotype;
  my $map;
  my $jaxmap;

  open(SNPMAP,$snpmap) || die "Could not open $snpmap\n";
  print "reading $snpmap\n";
  my $nsnp=0;
  $_ = <SNPMAP>;
  while (<SNPMAP>) {
    chomp;
    my($jax,@data) = split(/,/);
    $snpinfo->{$jax} = \@data;
    $genotype->{$jax}[0] = "  NA NA";
    $genotype->{$jax}[1] = "  $data[3] $data[3]";
    $genotype->{$jax}[2] = "  $data[3] $data[4]";
    $genotype->{$jax}[3] = "  $data[4] $data[4]";
    #    print "$jax $genotype->{$jax}[0] $genotype->{$jax}[1] $genotype->{$jax}[2] $genotype->{$jax}[3]\n";
    $igenotype->{$jax}[0] = undef;
    $igenotype->{$jax}[1] = $data[3].$data[3];
    $igenotype->{$jax}[3] = $data[4].$data[4];

    my $chr = $data[1];
    my $bp = $data[2];
    push @{$map->{$chr}}, [$jax, $bp];
    $jaxmap->{$jax} = [$chr, $bp];
    $nsnp++;
  }
  close(SNPMAP);
  print "read $nsnp snps\n";
  print "chrs " . join( " ", sort keys %$map) . "\n";
  return( $snpinfo, $genotype, $igenotype, $map, $jaxmap );

}

#_________________________________________________________________________________________
#
#  ReadInbredGenotypes
#  
#		Ignore non biallelic or with H or N
#		
#		$inbredgeno->{ProbeID}{founder} = genotype (A/C/G/T)
#
#_________________________________________________________________________________________
#
#	ProbeID       chr    Pos(build36)   posBuild37   rs.number    (129S1xCAST)F1   (129SIxPWK)F1   (129SIxPWK)F1.1   (129SIxPWK)F1.2   ...
#   JAX00725013   chry   37027          37027        rs46326766   N                N               N                 C                 ...
#   JAX00725015   chry   64107          64107        rs47485908   N                N               N                 A                 ...
sub ReadInbredGenotypes
{
	my ( $mdadir, $founders ) = @_;
	my $inbredgeno;
	my $ok;
	my $good;
	foreach my $file ( glob( "$mdadir/chr*.txt") ) 
	{
		#
		#	From header, map strain id to founder name
		#
		open(FILE, $file) || die "Could not open $file\n";
		$_ = <FILE>;
		chomp;
		my ( $ProbeID,$chr,$bp36,$bp37,$rs, @strain ) = split;
		my $use;
		foreach my $f ( @$founders )
		{
			for (my $k=0;$k<=$#strain;$k++) 
			{
				if ( $strain[$k] eq $f ) 
				{
					$use->{$k} = $f;
					#	  print "$file $f index $k\n"
				}
			}
		}

		while (<FILE>) 
		{
			chomp;
			my ( $ProbeID,$chr,$bp36,$bp37,$rs, @g ) = split;
			$ok = 1;
			#
			#	Only look at strains which map to founders.
			#	Note what bases are used (ACGT)
			#
			foreach my $k ( keys %$use ) 
			{
				if ( $g[$k] ne "H" && $g[$k] ne "N" ) 
				{
					$ac->{$ProbeID}{$g[$k]}++;
				} else 
				{
					$ok = 0;
				}
			}	
			#
			#	Only save when all ACGT (not N or H) and biallelic
			#
			my @a = keys %{$ac->{$ProbeID}};
			my $na = $#a+1;
			if ( $ok && $na ==2 ) 
			{
				$good++;
				foreach my $k ( keys %$use ) 
				{
					$inbredgeno->{$ProbeID}{$use->{$k}} = $g[$k];
				}	
			}	
		}	
	}	
	print "$good good snps in inbred data\n";
	return($inbredgeno, $ac);
}
